import Augmentor
import numpy as np

# 创建 Augmentor.Pipeline 对象
p = Augmentor.Pipeline("D:/work/data_augmentation")
# 添加旋转操作（随机旋转）
p.rotate_random_90(probability=1)
p.rotate_without_crop(probability=1, max_left_rotation=180, max_right_rotation=180)

# 添加缩放操作（随机缩放）
p.zoom_random(probability=1, percentage_area=0.8)

# 添加平移操作（随机平移）
p.random_distortion(probability=1, grid_width=4, grid_height=4, magnitude=8)

# 添加翻转操作（随机翻转）
p.flip_random(probability=1)

# 添加亮度调整操作（随机亮度调整）
p.random_brightness(probability=1, min_factor=0.8, max_factor=1.2)

# 添加对比度调整操作（随机对比度调整）
p.random_contrast(probability=1, min_factor=0.8, max_factor=1.2)


# # 添加噪声添加操作（随机噪声添加）
# class AddRandomNoise(Augmentor.Operations.Operation):
#     def __init__(self, probability):
#         super(AddRandomNoise, self).__init__(probability=probability)
#
#     def perform_operation(self, images):
#         for i in range(len(images)):
#             # 生成随机噪声
#             noise = np.random.normal(0, 1, images[i].shape)
#             # 将噪声添加到图像上
#             images[i] = np.clip(images[i] + noise, 0, 255).astype(np.uint8)
#         return images
#
#
# p.add_operation(AddRandomNoise(probability=1))

# 执行数据增强操作并生成增强后的图像
p.sample(5)  # 生成 100 张增强后的图像

# 可以在指定的输出目录中找到增强后的图像
